#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import plotly.graph_objects as go

# 设置支持中文的字体
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']


class Config:
    API_KEY = "1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c"  # 替换为您的 tushare API 密钥
    START_DATE = "20200101"
    END_DATE = "20221231"
    INITIAL_CAPITAL = 1_000_000
    STOP_LOSS = 0.05  # 止损比例
    TAKE_PROFIT = 0.10  # 止盈比例
    SHORT_WINDOW = 20  # 短期移动平均窗口
    LONG_WINDOW = 60  # 长期移动平均窗口
    TIME_FREQUENCY = "D"  # 时间粒度，选择：D（日线）、W（周线）、M（月线）、5（5分钟）、30（30分钟）、60（60分钟）
    TOP_N_STOCKS = 10  # 选股数量


class DataHandler:
    def __init__(self, config):
        self.config = config
        self.pro = ts.pro_api(self.config.API_KEY)

    def get_stock_basic(self):
        """获取股票列表"""
        return self.pro.stock_basic(fields='ts_code,symbol,name')

    def get_stock_data(self, ts_code, start_date, end_date, freq):
        """获取单支股票的历史数据"""
        if freq == "D":
            df = self.pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        elif freq == "W":
            df = self.pro.weekly(ts_code=ts_code, start_date=start_date, end_date=end_date)
        elif freq == "M":
            df = self.pro.monthly(ts_code=ts_code, start_date=start_date, end_date=end_date)
        elif freq in ["5", "30", "60"]:
            df = self.pro.min(ts_code=ts_code, start_date=start_date, end_date=end_date, freq=freq)
        else:
            print("不支持的时间粒度，请选择D、W、M、5、30、60。")
            return None
        if df.empty:
            print(f"股票 {ts_code} 在 {start_date} 到 {end_date} 之间没有数据。")
            return None
        return df

    def get_multi_stock_data(self, stock_list, start_date, end_date, freq):
        """获取多支股票的历史数据"""
        data_dict = {}
        for index, stock in stock_list.iterrows():
            ts_code = stock['ts_code']
            df = self.get_stock_data(ts_code, start_date, end_date, freq)
            if df is not None:
                data_dict[ts_code] = df
        if not data_dict:
            print("没有获取到有效的股票数据。")
            return None
        return data_dict


class Strategy:
    def __init__(self, config):
        self.config = config

    def moving_average_crossover(self, df):
        """均线交叉策略"""
        if df.empty:
            print("数据为空，无法应用策略。")
            return None
        df['short_ma'] = df['close'].rolling(window=self.config.SHORT_WINDOW).mean()
        df['long_ma'] = df['close'].rolling(window=self.config.LONG_WINDOW).mean()

        df['signal'] = 0
        buy_condition = (df['short_ma'] > df['long_ma']) & (df['short_ma'].shift(1) <= df['long_ma'].shift(1))
        sell_condition = (df['short_ma'] < df['long_ma']) & (df['short_ma'].shift(1) >= df['long_ma'].shift(1))
        df.loc[buy_condition, 'signal'] = 1
        df.loc[sell_condition, 'signal'] = -1

        if df['signal'].sum() == 0:
            print("没有生成任何交易信号，可能是数据不足或策略逻辑问题。")
            return None
        return df


class RiskManager:
    def __init__(self, config):
        self.config = config

    def apply_stop_loss(self, df):
        """止损机制"""
        if df.empty or 'close' not in df.columns:
            print("数据为空或没有'close'列，无法应用止损。")
            return df
        df['stop_loss_price'] = df['close'] * (1 - self.config.STOP_LOSS)
        return df

    def apply_take_profit(self, df):
        """止盈机制"""
        if df.empty or 'close' not in df.columns:
            print("数据为空或没有'close'列，无法应用止盈。")
            return df
        df['take_profit_price'] = df['close'] * (1 + self.config.TAKE_PROFIT)
        return df


class Backtester:
    def __init__(self, config):
        self.config = config

    def calculate_returns(self, df):
        """计算策略收益率"""
        if df.empty or 'signal' not in df.columns:
            print("数据为空或没有'signal'列，无法计算策略收益。")
            return None
        df['returns'] = df['close'].pct_change()
        df['strategy_returns'] = df['returns'] * df['signal'].shift(1)

        if df['strategy_returns'].isnull().all():
            print("'strategy_returns'列全为NaN，可能是因为信号列为空或计算错误。")
            return None
        return df

    def calculate_metrics(self, df):
        """计算回测指标"""
        try:
            if 'strategy_returns' not in df.columns or df['strategy_returns'].isnull().all():
                print("无法计算回测指标，'strategy_returns'列为空或全为NaN。")
                return None, None, None

            cumulative_returns = (1 + df['strategy_returns']).cumprod().iloc[-1]
            max_drawdown = (df['close'].cummax().cummin() / df['close'].cummax()).min()
            win_rate = (df['strategy_returns'] > 0).sum() / len(df['strategy_returns'])
            return cumulative_returns, max_drawdown, win_rate
        except Exception as e:
            print(f"计算回测指标时发生错误：{str(e)}")
            return None, None, None


class Visualizer:
    def __init__(self, config):
        self.config = config
        # 创建保存图片的目录
        self.image_dir = "images"
        if not os.path.exists(self.image_dir):
            os.makedirs(self.image_dir)

    def _save_plot(self, stock_name, plot_type):
        """保存当前的图表到指定的路径"""
        file_path = os.path.join(self.image_dir, f"{stock_name}_{plot_type}.png")
        try:
            plt.savefig(file_path, dpi=300, bbox_inches='tight')
            plt.close()
            print(f"图表已保存至：{file_path}")
        except Exception as e:
            print(f"保存图表时发生错误：{str(e)}")
            plt.close()

    def plot_k_line(self, df, stock_name):
        """绘制K线图并保存"""
        if df.empty or 'close' not in df.columns:
            print("数据为空或没有'close'列，无法绘制K线图。")
            return
        plt.figure(figsize=(16, 8))
        plt.plot(df['close'], label='收盘价')
        plt.title(f'{stock_name} K线图')
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True)
        plt.show()
        self._save_plot(stock_name, "k_line")

    def plot_strategy_signals(self, df, stock_name):
        """绘制交易信号图并保存"""
        if df.empty or 'signal' not in df.columns:
            print("数据为空或没有'signal'列，无法绘制交易信号。")
            return
        plt.figure(figsize=(16, 8))
        plt.plot(df['close'], label='收盘价')
        buy_signals = df[df['signal'] == 1]
        sell_signals = df[df['signal'] == -1]
        plt.scatter(buy_signals.index, buy_signals['close'],
                    label='买入信号', marker='^', color='g', s=100)
        plt.scatter(sell_signals.index, sell_signals['close'],
                    label='卖出信号', marker='v', color='r', s=100)
        plt.title(f'{stock_name} 交易信号')
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True)
        plt.show()
        self._save_plot(stock_name, "strategy_signals")

    def plot_returns_curve(self, df, stock_name):
        """绘制收益率曲线图并保存"""
        if df.empty or 'strategy_returns' not in df.columns:
            print("数据为空或没有'strategy_returns'列，无法绘制收益率曲线。")
            return
        plt.figure(figsize=(16, 8))
        plt.plot((1 + df['strategy_returns']).cumprod(), label='策略收益率')
        plt.title(f'{stock_name} 收益率曲线')
        plt.xlabel('日期')
        plt.ylabel('收益率')
        plt.legend()
        plt.grid(True)
        plt.show()
        self._save_plot(stock_name, "returns_curve")

    def plot_multi_stock_returns(self, stock_returns_dict):
        """绘制多支股票收益率曲线比较图并保存"""
        if not stock_returns_dict:
            print("没有多支股票的收益率数据进行比较。")
            return
        plt.figure(figsize=(16, 8))
        for ts_code, returns in stock_returns_dict.items():
            plt.plot(returns, label=ts_code)
        plt.title('多支股票收益率曲线比较')
        plt.xlabel('交易日')
        plt.ylabel('累计收益率')
        plt.legend()
        plt.grid(True)
        plt.show()
        file_path = os.path.join(self.image_dir, "multi_stock_returns.png")
        try:
            plt.savefig(file_path, dpi=300, bbox_inches='tight')
            plt.close()
            print(f"多支股票收益率曲线比较图已保存至：{file_path}")
        except Exception as e:
            print(f"保存多支股票收益率曲线比较图时发生错误：{str(e)}")
            plt.close()


def main():
    # 初始化配置
    config = Config()

    # 初始化数据处理
    data_handler = DataHandler(config)
    stock_list = data_handler.get_stock_basic()
    stock_list = stock_list.head(config.TOP_N_STOCKS)  # 选取前N只股票

    # 获取多支股票的数据
    data_dict = data_handler.get_multi_stock_data(stock_list, config.START_DATE, config.END_DATE, config.TIME_FREQUENCY)
    if data_dict is None or not data_dict:
        print("没有获取到有效的股票数据。")
        return

        # 初始化风险管理和回测器，并传递config
    risk_manager = RiskManager(config)
    backtester = Backtester(config)
    visualizer = Visualizer(config)

    # 存储回测结果
    stock_performance = {}

    # 遍历每只股票进行回测
    for ts_code in data_dict:
        print(f"\n回测股票：{ts_code}")
        df = data_dict[ts_code]

        # 应用策略
        strategy = Strategy(config)
        df_strategy = strategy.moving_average_crossover(df)
        if df_strategy is None:
            continue

            # 应用风险管理
        df_strategy = risk_manager.apply_stop_loss(df_strategy)
        df_strategy = risk_manager.apply_take_profit(df_strategy)

        # 计算策略收益
        df_backtest = backtester.calculate_returns(df_strategy)
        if df_backtest is None:
            continue

            # 计算回测指标
        cumulative_returns, max_drawdown, win_rate = backtester.calculate_metrics(df_backtest)
        if cumulative_returns is None:
            continue

            # 存储回测结果
        stock_performance[ts_code] = {
            '累计收益率': cumulative_returns,
            '最大回撤': max_drawdown,
            '胜率': win_rate
        }

        # 输出回测结果
        print(f"累计收益率: {cumulative_returns:.2%}")
        print(f"最大回撤: {max_drawdown:.2%}")
        print(f"胜率: {win_rate:.2%}")

        # 绘制可视化图表并保存
        visualizer.plot_k_line(df_strategy, ts_code)
        visualizer.plot_strategy_signals(df_strategy, ts_code)
        visualizer.plot_returns_curve(df_backtest, ts_code)

        # 绘制多支股票收益率曲线比较图并保存
    if stock_performance:
        stock_returns_dict = {}
        for ts_code in data_dict:
            df = data_dict[ts_code]
            strategy = Strategy(config)
            df_strategy = strategy.moving_average_crossover(df)
            if df_strategy is not None:
                df_backtest = backtester.calculate_returns(df_strategy)
                if df_backtest is not None and 'strategy_returns' in df_backtest.columns:
                    cumulative_returns = (1 + df_backtest['strategy_returns']).cumprod()
                    if not cumulative_returns.empty:
                        stock_returns_dict[ts_code] = cumulative_returns
        if stock_returns_dict:
            visualizer.plot_multi_stock_returns(stock_returns_dict)

            # 输出所有股票的回测结果
    print("\n所有股票的回测结果:")
    for ts_code, performance in stock_performance.items():
        print(f"股票代码: {ts_code}")
        print(f"累计收益率: {performance['累计收益率']:.2%}")
        print(f"最大回撤: {performance['最大回撤']:.2%}")
        print(f"胜率: {performance['胜率']:.2%}")
        print("-" * 50)


if __name__ == '__main__':
    main()

